{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c4eab109-d2c0-4409-936e-94f07050816e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "from peft import get_peft_model, LoraConfig,TaskType\n",
    "from peft.utils import prepare_model_for_kbit_training\n",
    "import torch\n",
    "from datasets import Dataset\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "37299903-fb62-49d9-ab37-a1d17ff74918",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_json('../datasets/cerbo.json')\n",
    "ds = Dataset.from_pandas(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "335f19ac-5320-4696-bbe2-5afd5210267c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['instruction', 'input', 'output'],\n",
       "    num_rows: 91\n",
       "})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5b841c2-115d-4d8d-ae5a-b7a315119e2b",
   "metadata": {},
   "source": [
    "## Processing the training dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "894c4ef3-36bb-4cd0-b50b-3b119296379d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "model_path = \"microsoft/Phi-3-mini-4k-instruct\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, padding_side=\"right\")\n",
    "tokenizer.pad_token = tokenizer.unk_token\n",
    "tokenizer.model_max_length = 2048\n",
    "tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e3055bab-07f7-470b-bbdc-fe777a4d4e19",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_chat_template(\n",
    "    example,\n",
    "    tokenizer,\n",
    "):\n",
    "    messages = [\n",
    "        {\"role\": \"user\", \"content\": example[\"instruction\"]},\n",
    "        {\"role\": \"assistant\", \"content\": example[\"output\"]},\n",
    "    ]\n",
    "    example[\"text\"] = tokenizer.apply_chat_template(\n",
    "        messages, tokenize=False, add_generation_prompt=False\n",
    "    )\n",
    "    return example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ea7b9679-c2a3-41a3-967e-f4073d639524",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d2f3f70cc42242f89334442fa5b4cb48",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Applying chat template to train_sft (num_proc=10):   0%|          | 0/91 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "processed_train_dataset = ds.map(\n",
    "        apply_chat_template,\n",
    "        fn_kwargs={\"tokenizer\": tokenizer},\n",
    "        num_proc=10,\n",
    "        remove_columns=ds.column_names,\n",
    "        desc=\"Applying chat template to train_sft\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8fe0cc3e-a5d4-4c1b-b436-ad162c52afad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'text': '<|user|>\\nhi<|end|>\\n<|assistant|>\\nHello! How can I assist you today?<|end|>\\n<|endoftext|>'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "processed_train_dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76017971-acca-43be-a51d-a96b8d54ae0e",
   "metadata": {},
   "source": [
    "## Create model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cd537edd-d959-4964-9c84-d4b4bbf7a196",
   "metadata": {},
   "outputs": [],
   "source": [
    "quantization_config = BitsAndBytesConfig(load_in_8bit=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9237ee9e-b956-44f0-b28a-b62556403aad",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`flash-attention` package not found, consider installing for better performance: No module named 'flash_attn'.\n",
      "Current `flash-attention` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "035a9060bf2c45b68bee146796e64375",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model = model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_path,\n",
    "        quantization_config=quantization_config,\n",
    "        device_map=\"auto\",\n",
    "        trust_remote_code=True,\n",
    "        torch_dtype=torch.bfloat16,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "eeb18420-554b-483b-8244-1186b4e358e2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Phi3ForCausalLM(\n",
       "  (model): Phi3Model(\n",
       "    (embed_tokens): Embedding(32064, 3072, padding_idx=32000)\n",
       "    (embed_dropout): Dropout(p=0.0, inplace=False)\n",
       "    (layers): ModuleList(\n",
       "      (0-31): 32 x Phi3DecoderLayer(\n",
       "        (self_attn): Phi3Attention(\n",
       "          (o_proj): Linear8bitLt(in_features=3072, out_features=3072, bias=False)\n",
       "          (qkv_proj): Linear8bitLt(in_features=3072, out_features=9216, bias=False)\n",
       "          (rotary_emb): Phi3RotaryEmbedding()\n",
       "        )\n",
       "        (mlp): Phi3MLP(\n",
       "          (gate_up_proj): Linear8bitLt(in_features=3072, out_features=16384, bias=False)\n",
       "          (down_proj): Linear8bitLt(in_features=8192, out_features=3072, bias=False)\n",
       "          (activation_fn): SiLU()\n",
       "        )\n",
       "        (input_layernorm): Phi3RMSNorm()\n",
       "        (resid_attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "        (resid_mlp_dropout): Dropout(p=0.0, inplace=False)\n",
       "        (post_attention_layernorm): Phi3RMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): Phi3RMSNorm()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=3072, out_features=32064, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "699e23ee-f508-4bf4-850f-5367abd61ee2",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.enable_input_require_grads()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e49a651-9de4-4911-beef-1dee87607828",
   "metadata": {},
   "source": [
    "## Lora"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3713dd42-7f7d-48cf-883a-8fb0a7f67a1a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LoraConfig(peft_type=<PeftType.LORA: 'LORA'>, auto_mapping=None, base_model_name_or_path=None, revision=None, task_type=<TaskType.CAUSAL_LM: 'CAUSAL_LM'>, inference_mode=False, r=8, target_modules={'gate_proj', 'o_proj', 'down_proj', 'v_proj', 'q_proj', 'up_proj', 'k_proj'}, lora_alpha=16, lora_dropout=0.055, fan_in_fan_out=False, bias='none', use_rslora=False, modules_to_save=None, init_lora_weights=True, layers_to_transform=None, layers_pattern=None, rank_pattern={}, alpha_pattern={}, megatron_config=None, megatron_core='megatron.core', loftq_config={}, use_dora=False, layer_replication=None)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM, \n",
    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
    "    inference_mode=False, \n",
    "    r=8, \n",
    "    lora_alpha=16, \n",
    "    lora_dropout=0.055,\n",
    "    bias=\"none\",\n",
    ")\n",
    "config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "950035e4-db83-4eda-82e4-2eff1b185cfa",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): LoraModel(\n",
       "    (model): Phi3ForCausalLM(\n",
       "      (model): Phi3Model(\n",
       "        (embed_tokens): Embedding(32064, 3072, padding_idx=32000)\n",
       "        (embed_dropout): Dropout(p=0.0, inplace=False)\n",
       "        (layers): ModuleList(\n",
       "          (0-31): 32 x Phi3DecoderLayer(\n",
       "            (self_attn): Phi3Attention(\n",
       "              (o_proj): lora.Linear8bitLt(\n",
       "                (base_layer): Linear8bitLt(in_features=3072, out_features=3072, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.055, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=3072, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=3072, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (qkv_proj): Linear8bitLt(in_features=3072, out_features=9216, bias=False)\n",
       "              (rotary_emb): Phi3RotaryEmbedding()\n",
       "            )\n",
       "            (mlp): Phi3MLP(\n",
       "              (gate_up_proj): Linear8bitLt(in_features=3072, out_features=16384, bias=False)\n",
       "              (down_proj): lora.Linear8bitLt(\n",
       "                (base_layer): Linear8bitLt(in_features=8192, out_features=3072, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.055, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=8192, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=3072, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (activation_fn): SiLU()\n",
       "            )\n",
       "            (input_layernorm): Phi3RMSNorm()\n",
       "            (resid_attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "            (resid_mlp_dropout): Dropout(p=0.0, inplace=False)\n",
       "            (post_attention_layernorm): Phi3RMSNorm()\n",
       "          )\n",
       "        )\n",
       "        (norm): Phi3RMSNorm()\n",
       "      )\n",
       "      (lm_head): Linear(in_features=3072, out_features=32064, bias=False)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = get_peft_model(model, config)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "aadefc23-66f2-4b32-af5d-4e22937ae052",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 4,456,448 || all params: 3,825,536,000 || trainable%: 0.1165\n"
     ]
    }
   ],
   "source": [
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3a0de08-9d0d-4ee2-8468-8eb719604334",
   "metadata": {},
   "source": [
    "## Configure training parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7f8fd582-4a27-4aca-a679-163f1efeb96e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import SFTConfig, SFTTrainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2c841522-7b9d-4a53-b4ba-af5e1d402d67",
   "metadata": {},
   "outputs": [],
   "source": [
    "args=SFTConfig(\n",
    "    output_dir=\"./output/Phi-3\",\n",
    "    per_device_train_batch_size=4,\n",
    "    gradient_accumulation_steps=4,\n",
    "    logging_steps=20,\n",
    "    log_level=\"info\",\n",
    "    num_train_epochs=50,\n",
    "    save_steps=100,\n",
    "    learning_rate=1e-4,\n",
    "    save_total_limit=2,\n",
    "    gradient_checkpointing=True,\n",
    "    dataset_text_field=\"text\",\n",
    "    max_seq_length=2048,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "be791275-244e-4cfe-a380-920abb06d513",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e66a893faaea4cafb6a5b34cb0b9a76e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/91 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from trl import SFTTrainer\n",
    "trainer = SFTTrainer(\n",
    "            model=model,\n",
    "            train_dataset=processed_train_dataset,\n",
    "            tokenizer=tokenizer,\n",
    "            args=args\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4afe388b-db7d-4073-8762-a3dc8c4eaaf2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "***** Running training *****\n",
      "  Num examples = 91\n",
      "  Num Epochs = 50\n",
      "  Instantaneous batch size per device = 4\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 16\n",
      "  Gradient Accumulation steps = 4\n",
      "  Total optimization steps = 250\n",
      "  Number of trainable parameters = 4,456,448\n",
      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n",
      "/home/zxd/miniconda3/envs/llm/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "You are not running the flash-attention implementation, expect numerical differences.\n",
      "/home/zxd/miniconda3/envs/llm/lib/python3.10/site-packages/bitsandbytes/autograd/_functions.py:316: UserWarning: MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='250' max='250' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [250/250 05:40, Epoch 43/50]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>1.649900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>40</td>\n",
       "      <td>0.844100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>60</td>\n",
       "      <td>0.652000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>80</td>\n",
       "      <td>0.556300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>0.460100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>120</td>\n",
       "      <td>0.374400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>140</td>\n",
       "      <td>0.297700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>160</td>\n",
       "      <td>0.231700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>180</td>\n",
       "      <td>0.188500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>0.159700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>220</td>\n",
       "      <td>0.142700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>240</td>\n",
       "      <td>0.129500</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to ./output/Phi-3/checkpoint-100\n",
      "/home/zxd/miniconda3/envs/llm/lib/python3.10/site-packages/peft/utils/save_and_load.py:195: UserWarning: Could not find a config file in /home/zxd/workspace/models/phi3-4k-mini - will assume that the vocabulary was not modified.\n",
      "  warnings.warn(\n",
      "tokenizer config file saved in ./output/Phi-3/checkpoint-100/tokenizer_config.json\n",
      "Special tokens file saved in ./output/Phi-3/checkpoint-100/special_tokens_map.json\n",
      "added tokens file saved in ./output/Phi-3/checkpoint-100/added_tokens.json\n",
      "/home/zxd/miniconda3/envs/llm/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/home/zxd/miniconda3/envs/llm/lib/python3.10/site-packages/bitsandbytes/autograd/_functions.py:316: UserWarning: MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n",
      "Saving model checkpoint to ./output/Phi-3/checkpoint-200\n",
      "/home/zxd/miniconda3/envs/llm/lib/python3.10/site-packages/peft/utils/save_and_load.py:195: UserWarning: Could not find a config file in /home/zxd/workspace/models/phi3-4k-mini - will assume that the vocabulary was not modified.\n",
      "  warnings.warn(\n",
      "tokenizer config file saved in ./output/Phi-3/checkpoint-200/tokenizer_config.json\n",
      "Special tokens file saved in ./output/Phi-3/checkpoint-200/special_tokens_map.json\n",
      "added tokens file saved in ./output/Phi-3/checkpoint-200/added_tokens.json\n",
      "/home/zxd/miniconda3/envs/llm/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/home/zxd/miniconda3/envs/llm/lib/python3.10/site-packages/bitsandbytes/autograd/_functions.py:316: UserWarning: MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n",
      "Saving model checkpoint to ./output/Phi-3/checkpoint-250\n",
      "/home/zxd/miniconda3/envs/llm/lib/python3.10/site-packages/peft/utils/save_and_load.py:195: UserWarning: Could not find a config file in /home/zxd/workspace/models/phi3-4k-mini - will assume that the vocabulary was not modified.\n",
      "  warnings.warn(\n",
      "tokenizer config file saved in ./output/Phi-3/checkpoint-250/tokenizer_config.json\n",
      "Special tokens file saved in ./output/Phi-3/checkpoint-250/special_tokens_map.json\n",
      "added tokens file saved in ./output/Phi-3/checkpoint-250/added_tokens.json\n",
      "Deleting older checkpoint [output/Phi-3/checkpoint-100] due to args.save_total_limit\n",
      "\n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=250, training_loss=0.46030655479431154, metrics={'train_runtime': 343.1356, 'train_samples_per_second': 13.26, 'train_steps_per_second': 0.729, 'total_flos': 9045626707445760.0, 'train_loss': 0.46030655479431154, 'epoch': 43.47826086956522})"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a44e45b-ce94-4e99-8ce8-54e32df898a3",
   "metadata": {},
   "source": [
    "## Save LoRA and tokenizer results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "1254dd7c-4602-4f2c-a702-042043a47ce1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/zxd/miniconda3/envs/llm/lib/python3.10/site-packages/peft/utils/save_and_load.py:195: UserWarning: Could not find a config file in /home/zxd/workspace/models/phi3-4k-mini - will assume that the vocabulary was not modified.\n",
      "  warnings.warn(\n",
      "tokenizer config file saved in ./Phi-3_lora/tokenizer_config.json\n",
      "Special tokens file saved in ./Phi-3_lora/special_tokens_map.json\n",
      "added tokens file saved in ./Phi-3_lora/added_tokens.json\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "('./Phi-3_lora/tokenizer_config.json',\n",
       " './Phi-3_lora/special_tokens_map.json',\n",
       " './Phi-3_lora/tokenizer.model',\n",
       " './Phi-3_lora/added_tokens.json')"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lora_path='./Phi-3_lora'\n",
    "trainer.model.save_pretrained(lora_path)\n",
    "tokenizer.save_pretrained(lora_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cea5a92-f9fc-4229-ac2e-1d92fb26e33d",
   "metadata": {},
   "source": [
    "## Load lora weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ea0d76ee-c0eb-438d-ae0f-61b02fddd3e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`flash-attention` package not found, consider installing for better performance: No module named 'flash_attn'.\n",
      "Current `flash-attention` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e0fbd7f33a5d42c1b1285262f3d6d8ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "import torch \n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline \n",
    "import torch\n",
    "from peft import (\n",
    "    PeftModel,\n",
    "    LoraConfig,\n",
    "    get_peft_model,\n",
    "    get_peft_model_state_dict,\n",
    "    prepare_model_for_kbit_training,\n",
    "    set_peft_model_state_dict,\n",
    ")\n",
    "model_path = \"microsoft/Phi-3-mini-4k-instruct\"\n",
    "torch.random.manual_seed(0) \n",
    "model = AutoModelForCausalLM.from_pretrained( \n",
    "    model_path,  \n",
    "    device_map=\"cuda\",  \n",
    "    torch_dtype=\"auto\",  \n",
    "    trust_remote_code=True,  \n",
    ")\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "266bb4e1-eed5-4fd8-9b93-b797404099c8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/zxd/miniconda3/envs/llm/lib/python3.10/site-packages/transformers/generation/configuration_utils.py:540: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.0` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`.\n",
      "  warnings.warn(\n",
      "The `seen_tokens` attribute is deprecated and will be removed in v4.41. Use the `cache_position` model input instead.\n",
      "You are not running the flash-attention implementation, expect numerical differences.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " I am Phi, Microsoft's language model, trained to assist with a wide range of queries and tasks.\n"
     ]
    }
   ],
   "source": [
    "prompt = \"who are you?\"\n",
    "messages = [ \n",
    "    {\"role\": \"user\", \"content\":prompt}, \n",
    "] \n",
    "\n",
    "pipe = pipeline( \n",
    "    \"text-generation\", \n",
    "    model=model, \n",
    "    tokenizer=tokenizer, \n",
    ") \n",
    "\n",
    "generation_args = { \n",
    "    \"max_new_tokens\": 500, \n",
    "    \"return_full_text\": False, \n",
    "    \"temperature\": 0.0, \n",
    "    \"do_sample\": False, \n",
    "} \n",
    "\n",
    "output = pipe(messages, **generation_args) \n",
    "print(output[0]['generated_text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fbaa84c7-1cd6-4c96-be54-2b4e7b1489d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "lora_config_path = './Phi-3_lora'\n",
    "config = LoraConfig.from_pretrained(lora_config_path)\n",
    "\n",
    "model = PeftModel.from_pretrained(model, model_id=lora_config_path, config=config)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f5195288-5e11-4c86-ba14-79f208169a56",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): LoraModel(\n",
       "    (model): Phi3ForCausalLM(\n",
       "      (model): Phi3Model(\n",
       "        (embed_tokens): Embedding(32064, 3072, padding_idx=32000)\n",
       "        (embed_dropout): Dropout(p=0.0, inplace=False)\n",
       "        (layers): ModuleList(\n",
       "          (0-31): 32 x Phi3DecoderLayer(\n",
       "            (self_attn): Phi3Attention(\n",
       "              (o_proj): lora.Linear(\n",
       "                (base_layer): Linear(in_features=3072, out_features=3072, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.055, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=3072, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=3072, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (qkv_proj): Linear(in_features=3072, out_features=9216, bias=False)\n",
       "              (rotary_emb): Phi3RotaryEmbedding()\n",
       "            )\n",
       "            (mlp): Phi3MLP(\n",
       "              (gate_up_proj): Linear(in_features=3072, out_features=16384, bias=False)\n",
       "              (down_proj): lora.Linear(\n",
       "                (base_layer): Linear(in_features=8192, out_features=3072, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.055, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=8192, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=3072, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (activation_fn): SiLU()\n",
       "            )\n",
       "            (input_layernorm): Phi3RMSNorm()\n",
       "            (resid_attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "            (resid_mlp_dropout): Dropout(p=0.0, inplace=False)\n",
       "            (post_attention_layernorm): Phi3RMSNorm()\n",
       "          )\n",
       "        )\n",
       "        (norm): Phi3RMSNorm()\n",
       "      )\n",
       "      (lm_head): Linear(in_features=3072, out_features=32064, bias=False)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b0b58b71-3766-463a-b4be-87bc486712f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The model 'PeftModelForCausalLM' is not supported for text-generation. Supported models are ['BartForCausalLM', 'BertLMHeadModel', 'BertGenerationDecoder', 'BigBirdForCausalLM', 'BigBirdPegasusForCausalLM', 'BioGptForCausalLM', 'BlenderbotForCausalLM', 'BlenderbotSmallForCausalLM', 'BloomForCausalLM', 'CamembertForCausalLM', 'LlamaForCausalLM', 'CodeGenForCausalLM', 'CohereForCausalLM', 'CpmAntForCausalLM', 'CTRLLMHeadModel', 'Data2VecTextForCausalLM', 'DbrxForCausalLM', 'ElectraForCausalLM', 'ErnieForCausalLM', 'FalconForCausalLM', 'FuyuForCausalLM', 'GemmaForCausalLM', 'Gemma2ForCausalLM', 'GitForCausalLM', 'GPT2LMHeadModel', 'GPT2LMHeadModel', 'GPTBigCodeForCausalLM', 'GPTNeoForCausalLM', 'GPTNeoXForCausalLM', 'GPTNeoXJapaneseForCausalLM', 'GPTJForCausalLM', 'JambaForCausalLM', 'JetMoeForCausalLM', 'LlamaForCausalLM', 'MambaForCausalLM', 'MarianForCausalLM', 'MBartForCausalLM', 'MegaForCausalLM', 'MegatronBertForCausalLM', 'MistralForCausalLM', 'MixtralForCausalLM', 'MptForCausalLM', 'MusicgenForCausalLM', 'MusicgenMelodyForCausalLM', 'MvpForCausalLM', 'OlmoForCausalLM', 'OpenLlamaForCausalLM', 'OpenAIGPTLMHeadModel', 'OPTForCausalLM', 'PegasusForCausalLM', 'PersimmonForCausalLM', 'PhiForCausalLM', 'Phi3ForCausalLM', 'PLBartForCausalLM', 'ProphetNetForCausalLM', 'QDQBertLMHeadModel', 'Qwen2ForCausalLM', 'Qwen2MoeForCausalLM', 'RecurrentGemmaForCausalLM', 'ReformerModelWithLMHead', 'RemBertForCausalLM', 'RobertaForCausalLM', 'RobertaPreLayerNormForCausalLM', 'RoCBertForCausalLM', 'RoFormerForCausalLM', 'RwkvForCausalLM', 'Speech2Text2ForCausalLM', 'StableLmForCausalLM', 'Starcoder2ForCausalLM', 'TransfoXLLMHeadModel', 'TrOCRForCausalLM', 'WhisperForCausalLM', 'XGLMForCausalLM', 'XLMWithLMHeadModel', 'XLMProphetNetForCausalLM', 'XLMRobertaForCausalLM', 'XLMRobertaXLForCausalLM', 'XLNetLMHeadModel', 'XmodForCausalLM'].\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I am Fairy, an AI assistant developed by Cerbo AI.\n"
     ]
    }
   ],
   "source": [
    "messages = [ \n",
    "    {\"role\": \"user\", \"content\":prompt}, \n",
    "] \n",
    "\n",
    "pipe = pipeline( \n",
    "    \"text-generation\", \n",
    "    model=model, \n",
    "    tokenizer=tokenizer, \n",
    ") \n",
    "\n",
    "generation_args = { \n",
    "    \"max_new_tokens\": 500, \n",
    "    \"return_full_text\": False, \n",
    "    \"temperature\": 0.0, \n",
    "    \"do_sample\": False, \n",
    "} \n",
    "\n",
    "output = pipe(messages, **generation_args) \n",
    "print(output[0]['generated_text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "557ca88f-8961-4140-a796-8e1181faf6af",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
